import os
from PIL import Image
import torchvision.transforms as transforms
import argparse

img2tensor = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]) 

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--dir_path', default='/data1/2021/lyf/GcGAN/results/SYN_d2n_rot/test_latest/images')
    args = parser.parse_args()
    return args

def calc_mean_std(feat, eps=1e-5):
    # eps is a small value added to the variance to avoid divide-by-zero.
    size = feat.size()
    assert (len(size) == 4)
    N, C = size[:2]
    feat_var = feat.view(N, C, -1).var(dim=2) + eps
    feat_std = feat_var.sqrt().view(N, C, 1, 1)
    feat_mean = feat.view(N, C, -1).mean(dim=2).view(N, C, 1, 1)
    return feat_mean, feat_std

def calc_img_mean_std(img_path):   
    x1=Image.open(img_path)
    x1=img2tensor(x1).unsqueeze(0)
    mean, std=calc_mean_std(x1)
    return mean, std


def calc_dir_mean_std(dir_path):
    files=os.listdir(dir_path)
    mean, std, len_files = 0, 0, 0
    for file in files:
        img_file = os.path.join(dir_path, file)
        img_mean, img_std = calc_img_mean_std(img_file)
        mean += img_mean
        std += img_std
        len_files += 1
    print("# of pair images:", len_files)
    mean /= len_files
    std /= len_files
    return mean, std

if __name__ == '__main__':
    args = parse_args()
    # path1 = 
    mean, std = calc_dir_mean_std(args.dir_path)
    print("mean = " , mean)
    print("std = " , std)